from construct2 import *
from layer2.ethernet import ethernet_header
from layer3.ipv4 import ipv4_header
from layer3.ipv6 import ipv6_header
from layer4.tcp import tcp_header
from layer4.udp import udp_header


layer4_tcp = Struct("layer4_tcp",
    Rename("header", tcp_header),
    HexDumpAdapter(
        MetaField("next", lambda ctx: 
            ctx["_"]["header"].payload_length - ctx["header"].header_length
        )
    ),
)

layer4_udp = Struct("layer4_udp",
    Rename("header", udp_header),
    HexDumpAdapter(
        MetaField("udp_raw", lambda ctx: ctx["header"].payload_length)
    ),    
)

layer3_ipv4 = Struct("layer3_ipv4",
    Rename("header", ipv4_header),
    Switch("next", lambda ctx: ctx["header"].protocol,
        {
            "TCP" : layer4_tcp,
            "UDP" : layer4_udp,
        },
        default = Pass
    )
)

layer2_ethernet = Struct("layer2_ethernet",
    Rename("header", ethernet_header),
    Switch("next", lambda ctx: ctx["header"].type,
        {
            "IPv4" : layer3_ipv4,
        },
        default = Pass,
    )
)

ip_stack = Rename("ip_stack", layer2_ethernet)


if __name__ == "__main__":
    cap1 = (
    "0011508c283c001150886b570800450001e971474000800684e4c0a80202525eedda112a0050d98e"
    "c61d54fe977d501844705dcc0000474554202f20485454502f312e310d0a486f73743a207777772e"
    "707974686f6e2e6f72670d0a557365722d4167656e743a204d6f7a696c6c612f352e30202857696e"
    "646f77733b20553b2057696e646f7773204e5420352e313b20656e2d55533b2072763a312e382e30"
    "2e3129204765636b6f2f32303036303131312046697265666f782f312e352e302e310d0a41636365"
    "70743a20746578742f786d6c2c6170706c69636174696f6e2f786d6c2c6170706c69636174696f6e"
    "2f7868746d6c2b786d6c2c746578742f68746d6c3b713d302e392c746578742f706c61696e3b713d"
    "302e382c696d6167652f706e672c2a2f2a3b713d302e350d0a4163636570742d4c616e6775616765"
    "3a20656e2d75732c656e3b713d302e350d0a4163636570742d456e636f64696e673a20677a69702c"
    "6465666c6174650d0a4163636570742d436861727365743a2049534f2d383835392d312c7574662d"
    "383b713d302e372c2a3b713d302e370d0a4b6565702d416c6976653a203330300d0a436f6e6e6563"
    "74696f6e3a206b6565702d616c6976650d0a507261676d613a206e6f2d63616368650d0a43616368"
    "652d436f6e74726f6c3a206e6f2d63616368650d0a0d0a"
    ).decode("hex")
    
    cap2 = (
    "0011508c283c0002e342600908004500038ea0e8400080069b82c0a8020540e9b7680d"
    "b5005062303fb21836e9e650184470c9bc0000474554202f20485454502f312e310d0a"
    "486f73743a207777772e676f6f676c652e636f6d0d0a557365722d4167656e743a204d"
    "6f7a696c6c612f352e30202857696e646f77733b20553b2057696e646f7773204e5420"
    "352e313b20656e2d55533b2072763a312e382e302e3629204765636b6f2f3230303630"
    "3732382046697265666f782f312e352e302e360d0a4163636570743a20746578742f78"
    "6d6c2c6170706c69636174696f6e2f786d6c2c6170706c69636174696f6e2f7868746d"
    "6c2b786d6c2c746578742f68746d6c3b713d302e392c746578742f706c61696e3b713d"
    "302e382c696d6167652f706e672c2a2f2a3b713d302e350d0a4163636570742d4c616e"
    "67756167653a20656e2d75732c656e3b713d302e350d0a4163636570742d456e636f64"
    "696e673a20677a69702c6465666c6174650d0a4163636570742d436861727365743a20"
    "49534f2d383835392d312c7574662d383b713d302e372c2a3b713d302e370d0a4b6565"
    "702d416c6976653a203330300d0a436f6e6e656374696f6e3a206b6565702d616c6976"
    "650d0a436f6f6b69653a2072656d656d6265726d653d747275653b204750433d46573d"
    "313a4748563d303a53473d303a54453d303a54533d303a54563d303a5349473d514561"
    "484c514c58785353624f765f5f3b20505245463d49443d333962383364666136653465"
    "346430313a46463d343a4c443d656e3a4e523d33303a43523d323a544d3d3131353533"
    "31333030313a4c4d3d313135353939383734303a474d3d313a533d4449424239745a57"
    "6543796b504530373b205349443d44514141414863414141414b6977492d775061784c"
    "6d7154744a4d3754685a417a365f724973356941327766476343784b71447432526949"
    "695366597666457676733565453073517a546c4e464e4b556f4f2d5169526e6b4e3774"
    "76644b633965586277704e6356334d67764e2d4a46374648795f4f435044324a595945"
    "50787a727574546e6972673434787a434a6373706d624e4c7a7879636b4e7858685f49"
    "5a516b63544a5a3256336573304b473671797a4f673b20533d676d61696c3d56756631"
    "7632512d754a453a676d61696c5f796a3d67477544523543495132633a676d70726f78"
    "793d75746467644636513142773a676d70726f78795f796a3d672d34417a69746e4343"
    "6f3b20545a3d2d3132300d0a0d0a"
    ).decode("hex")
    
#    obj = ip_stack.parse(cap1)
#    print obj
#    print repr(ip_stack.build(obj))
#    
#    print "-" * 80
#    
#    obj = ip_stack.parse(cap2)
#    print obj
#    print repr(ip_stack.build(obj))
    
    import time
    from Construct.Inventory.IpStack import ip_stack as old_ip_stack, seq_ip_stack
    new_obj1 = ip_stack.parse(cap1)
    old_obj1 = old_ip_stack.parse(cap1)
    new_obj2 = ip_stack.parse(cap2)
    old_obj2 = old_ip_stack.parse(cap2)
    
    def timeit(func, arg, count = 5000):
        t = time.time()
        while count:
            func(arg)
            count -= 1
        return time.time() - t
    
    print "parse cap1"
    print "new", timeit(ip_stack.parse, cap1)
    print "old", timeit(old_ip_stack.parse, cap1)
    print
    
    print "parse cap2"
    print "new", timeit(ip_stack.parse, cap2)
    print "old", timeit(old_ip_stack.parse, cap2)
    print
    
    print "build cap1"
    print "new", timeit(ip_stack.build, new_obj1)
    print "old", timeit(old_ip_stack.build, old_obj1)
    print
    
    print "build cap2"
    print "new", timeit(ip_stack.build, new_obj2)
    print "old", timeit(old_ip_stack.build, old_obj2)
    print
    
    #parse cap1 (small)
    #new 4.68700003624
    #old 6.67200016975 
    # --> 42% faster
    #
    #parse cap2 (big)
    #new 4.43799996376
    #old 9.07799983025
    # --> 104% faster
    #
    #build cap1 (small)
    #new 2.96900010109
    #old 7.57800006866
    # --> 155% faster
    #
    #build cap2 (big)
    #new 2.96799993515
    #old 9.96900010109
    # --> 236% faster










